#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Nov 12 2021
Plot data calculated from offline dry deposition code for South America 
@author: arifeinberg
"""

import matplotlib.pyplot as plt
import numpy as np
import cartopy.crs as ccrs
import scipy.io as sio
import datetime
import xarray as xr

#%% load GC dry deposition data
pn1 = '../GEOS-Chem_runs/run0019/OutputDir/'
pn2 = '../GEOS-Chem_runs/run0020/OutputDir/'
fn = 'GEOSChem.DryDep.alltime_m.nc4'

#Load all datasets
ds1 = xr.open_dataset(pn1 + fn)
ds2 = xr.open_dataset(pn2 + fn)

# dry deposition velocities in cm s^-1
Hg0_vel_R = ds1.DryDepVel_Hg0.isel(time=slice(12,24)).mean("time") # monthly
Hg0_vel_S = ds2.DryDepVel_Hg0.isel(time=slice(12,24)).mean("time") # monthly

#%% load data file

# get latitude and longitude for amazon area
fn_ols2 = 'misc_Data/Olson_2001_Land_Type_Masks.2_25.generic.nc'
ds_ols2 = xr.open_dataset(fn_ols2)
Olson_landtype = ds_ols2.to_array().squeeze()

fn_ols1 = 'misc_Data/Olson_2001_Drydep_Inputs.nc'
ds_ols1 = xr.open_dataset(fn_ols1)

lon = ds_ols2.lon # longitude
lat = ds_ols2.lat # latitude

# select Amazon coordinate bounds
lat_min = -34
lat_max = 14
lon_min = -82
lon_max = -33

# restrict lat and lon
lat_A = lat[(lat<=lat_max) & (lat>=lat_min)]
lon_A = lon[(lon<=lon_max) & (lon>=lon_min)]

#%% restrict to amazon area the Olson landtype
Olson_landtype_A = Olson_landtype[:,(lat<=lat_max)&(lat>=lat_min),(lon<=lon_max)&(lon>=lon_min)]
Hg0_vel_S_A = Hg0_vel_S[(lat<=lat_max)&(lat>=lat_min),(lon<=lon_max)&(lon>=lon_min)]
Hg0_vel_R_A = Hg0_vel_R[(lat<=lat_max)&(lat>=lat_min),(lon<=lon_max)&(lon>=lon_min)]

# create array of Amazon rainforest indices
IDEP = ds_ols1.IDEP.values # Mapping index: Olson land type ID to drydep ID
IOLSON_6 = np.asarray(np.where(IDEP==6)).flatten()

ind_rainforest = np.asarray(np.where((sum(Olson_landtype_A[IOLSON_6,:,:]) > 0.1)))

Hg0_vel_S_v = Hg0_vel_S_A.values
Hg0_vel_R_v = Hg0_vel_R_A.values

dep_vel_median_S = np.median(Hg0_vel_S_v[ind_rainforest[0,:],ind_rainforest[1,:]])
dep_vel_median_R = np.median(Hg0_vel_R_v[ind_rainforest[0,:],ind_rainforest[1,:]])

f, axes = plt.subplots(1, 1, figsize=[12,8],subplot_kw=dict(projection=ccrs.PlateCarree()),
                              gridspec_kw=dict(hspace=0.4, wspace=0.1))
h = Hg0_vel_R.plot.pcolormesh("lon", "lat", vmin=0,vmax=0.4, add_colorbar=False,
                               rasterized = True)
cbar = plt.colorbar(h, ax=axes, fraction=0.046, pad=0.04,extend='max',
                    orientation='horizontal')
cbar.set_label('Hg$^0$ dry dep velocity (cm s$^{-1}$)', fontsize=30)   
cbar.ax.tick_params(labelsize=28)
axes.set_xlim([lon_min, lon_max])
axes.set_ylim([lat_min, lat_max])
axes.coastlines(linewidth=2)
f.savefig('Figures/Fig5a_RF_DV_GC.pdf',bbox_inches = 'tight')

f, axes = plt.subplots(1, 1, figsize=[12,8],subplot_kw=dict(projection=ccrs.PlateCarree()),
                              gridspec_kw=dict(hspace=0.4, wspace=0.1))
h = Hg0_vel_S.plot.pcolormesh("lon", "lat", vmin=0,vmax=0.4, add_colorbar=False,
                               rasterized = True)
cbar = plt.colorbar(h, ax=axes, fraction=0.046, pad=0.04,extend='max',
                    orientation='horizontal')
cbar.set_label('Hg$^0$ dry dep velocity (cm s$^{-1}$)', fontsize=30)   
cbar.ax.tick_params(labelsize=28)
axes.set_xlim([lon_min, lon_max])
axes.set_ylim([lat_min, lat_max])
axes.coastlines(linewidth=2)
f.savefig('Figures/Fig5a_SV_DV_GC.pdf',bbox_inches = 'tight')
